import cv2
from ultralytics import YOLO

# 初始化视频捕获
cap = cv2.VideoCapture('2.mp4')

# 加载 YOLOv8 模型
model = YOLO('yolov8n.pt')

# 获取视频帧率，用于计算速度
fps = cap.get(cv2.CAP_PROP_FPS)

# 定义拥堵速度阈值（可根据实际情况调整），单位：像素/帧
congestion_speed_threshold = 5

# 存储上一帧车辆的信息，键为车辆 ID，值为 (上一帧位置, 上一帧时间)
vehicle_info = {}
next_vehicle_id = 0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 获取当前帧的时间
    current_time = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000

    # 使用 YOLOv8 进行目标检测
    results = model(frame, classes=2)  # 假设车辆类别编号为 2

    # 存储当前帧检测到的车辆信息，键为车辆 ID，值为 (当前位置, 当前时间)
    current_vehicles = {}

    for result in results:
        boxes = result.boxes.cpu().numpy()
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            current_position = ((x1 + x2) / 2, (y1 + y2) / 2)  # 取中心点作为车辆位置

            # 尝试匹配已有的车辆
            matched = False
            for vehicle_id, (prev_position, prev_time) in vehicle_info.items():
                prev_x, prev_y = prev_position
                dx = abs(current_position[0] - prev_x)
                dy = abs(current_position[1] - prev_y)
                distance = (dx ** 2 + dy ** 2) ** 0.5
                time_elapsed = current_time - prev_time
                if time_elapsed > 0:
                    speed = distance / time_elapsed
                else:
                    speed = 0

                # 如果距离变化小，认为是同一辆车
                if distance < 50:
                    # 匹配成功，更新车辆信息
                    current_vehicles[vehicle_id] = (current_position, current_time)
                    vehicle_info[vehicle_id] = (current_position, current_time)
                    matched = True
                    break

            if not matched:
                # 新车辆，分配新的 ID
                vehicle_id = next_vehicle_id
                next_vehicle_id += 1
                current_vehicles[vehicle_id] = (current_position, current_time)
                vehicle_info[vehicle_id] = (current_position, current_time)

            # 绘制车辆矩形框
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 移除当前帧未检测到的车辆
    vehicles_to_remove = []
    for vehicle_id in vehicle_info:
        if vehicle_id not in current_vehicles:
            vehicles_to_remove.append(vehicle_id)
    for vehicle_id in vehicles_to_remove:
        del vehicle_info[vehicle_id]

    # 检查是否有车辆速度低于阈值
    congested = False
    for vehicle_id, (prev_position, prev_time) in vehicle_info.items():
        current_position, current_time = current_vehicles.get(vehicle_id, (None, None))
        if current_position is not None:
            prev_x, prev_y = prev_position
            dx = abs(current_position[0] - prev_x)
            dy = abs(current_position[1] - prev_y)
            distance = (dx ** 2 + dy ** 2) ** 0.5
            time_elapsed = current_time - prev_time
            if time_elapsed > 0:
                speed = distance / time_elapsed
            else:
                speed = 0
            if speed < congestion_speed_threshold:
                congested = True
                break

    # 预测拥堵
    if congested:
        cv2.putText(frame, "Congestion Detected!", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    else:
        cv2.putText(frame, "Normal Traffic", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

    # 显示结果
    cv2.imshow('Traffic Congestion Prediction', frame)

    # 按 'q' 键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()
